from torch.utils import data 
import numpy as np
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda, Compose
import matplotlib.pyplot as plt

# 1. 从内存中加载数据并转换为tensor格式的数据
# 2. 根据非独立同分布的设置，讲数据划分为多分，并返回多分数据的numpy数据数据。因为返回dataset反而不利于进一步的划分和处理。
# 主要完成数据的多个客户端划分。不划分训练数据还是测试数据。因为可能存在一下多种情况：
# a. 分为测试节点和训练节点
# b. 每个客户端既包含测试数据也包含训练数据。单个客户端还可能面临其他情况的划分。
class DataManager():
    def __init__(self,client_number=20):
        # 客户端的数量
        self.client_number = int(client_number);
        

    # 1 加载数据、并对数据进行预处理。转换为Tensor格式
    def load_data(self):
        train_dataset = datasets.MNIST(
            '../../data', train=True, transform=ToTensor(), download=True
        )
        test_dataset = datasets.MNIST(
            '../../data', train=False, transform=ToTensor(),download=True
        )

        # self.train_data = np.vstack((train_dataset.data.numpy(),test_dataset.data.numpy()))
        # self.train_targets =np.hstack((train_dataset.targets.numpy(),test_dataset.targets.numpy()))
        self.train_data = train_dataset.data.numpy()
        self.train_targets =train_dataset.targets.numpy()
        self.test_data = train_dataset.data.numpy()
        self.test_targets =train_dataset.targets.numpy()
        print("特征维度:",self.train_data.shape,self.train_data.dtype)
        print("标签维度:",self.train_targets.shape,self.train_targets.dtype)
    

    # 根据客户端的总量互斥平均，分配成client_number份。假定原来的数据已经shuffle了。
    def allocate_data_avg(self):
        data_size=len(self.train_data)
        data_counts = int(data_size/self.client_number)
        dataset_list = []

        for i in range(self.client_number):
            # print(i*data_counts,(i+1)*data_counts)
            dataset_list.append((self.train_data[i*data_counts:(i+1)*data_counts],self.train_targets[i*data_counts:(i+1)*data_counts]))

        print("avg----client_size:",self.client_number,"----data_counts:",data_counts)
        return dataset_list

    # 从数据集中随机选取固定数量的数据。
    # 基于随机选择random.choice和numpy数组的花式索引实现。
    # data_counts表示每个客户端具有的数据规模
    def allocate_data_iid(self,data_counts=300):
        # 数据随机选择
        dataset_list = []
        for i in range(self.client_number):
            idx = np.random.choice(range(len(self.train_data)),data_counts,replace=False)
            dataset_list.append((self.train_data[idx],self.train_targets[idx]))
        
        print("iid----client_size:",self.client_number,"----data_counts:",data_counts)
        return dataset_list

    # 从数据集中选取随机类别的固定数量的数据。一种典型的非独立同分布划分方法
    # 基于随机选择random.choice和numpy数组的花式索引实现。
    # class_counts 表示类别的数量
    # data_counts 表示数据的规模
    def allocate_data_noniid(self,class_total=10,class_counts=2,data_counts=300):
        # 数据随机选择
        dataset_list = []
        
        # 每个客户端随机选择标签
        label_list = [np.random.choice(range(class_total),class_counts,replace=False) for i in range(self.client_number)]

        # 每个类别的数据放到一块
        data_class = []
        target_class = []
        for i in range(class_total):
            data_class.append(self.train_data[self.train_targets==i])
            target_class.append(np.ones(len(self.train_data[self.train_targets==i]),dtype=np.int64)*i)
        # 没办法直接转换成二维数组，因为第二个维度不一致
        # 给每个客户端选择指定数量数据，限制在某一类别。从lable_list确定的3个类中，选择数据。
        for labels in label_list:
            class_cat = []
            target_cat = []
            for i in labels:
                class_cat.append(data_class[i])
                target_cat.append(target_class[i])

            
            class_cat=np.vstack(class_cat)
            target_cat=np.hstack(target_cat)
 
            size = len(class_cat)
            idx = np.random.choice(range(size),data_counts,replace=False)
            dataset_list.append((class_cat[idx],target_cat[idx]))
        
        print("noniid----client_size:",self.client_number,"----data_counts:",data_counts,"----class_counts:",class_counts)
        return dataset_list

    def allocate_data_test(self):
        return [(self.test_data,self.test_targets),]
    # 使用TL散度计算当前数据非独立同分布的程度（或者其他方式，用来衡量划分后的数据的非独立同分布的程度）
    def calculate_degree(self,dataset_list):
        pass 